import os
import unittest

from ascend_deployer.trans_excel_to_inventory import HostInfo, HccnHostInfo, CommonHostInfo, InventoryInfo


class TestTransExcelToInventory(unittest.TestCase):
    def setUp(self):
        table_header = ["Title【勿动】", "IP", "ansible_ssh_user", "ansible_ssh_pass", "ansible_ssh_port", "set_hostname",
                        "npu_num", "davinci", "NPU0", "NPU1", "NPU2","NPU3","NPU4",	"NPU5", "NPU6", "NPU7", "NPU8",
                        "NPU9", "NPU10", "NPU11", "NPU12", "NPU13", "NPU14", "NPU15", "NPU0", "NPU1", "NPU2", "NPU3",
                        "NPU4", "NPU5", "NPU6", "NPU7", "NPU8", "NPU9", "NPU10", "NPU11", "NPU12", "NPU13", "NPU14",
                        "NPU15", "gateways", "netmask", "bitmap", "dscp_tc", "common_network", "roce_port", "SCALE",
                        "RUNNER_IP", "WEIGHTS_PATH"]
        self.host_info = HostInfo([], table_header)

    def test_get_option_value(self):
        ip = "1.1.1.1"
        user = "root"
        password = "test"
        data_list = [ip, user, password]
        table_header = ["IP", "ansible_ssh_user", "ansible_ssh_pass"]
        self.assertEqual(ip, self.host_info.get_option_value("IP", data_list, table_header))
        self.assertEqual(user, self.host_info.get_option_value("ansible_ssh_user", data_list, table_header))
        self.assertEqual(password, self.host_info.get_option_value("ansible_ssh_pass",
                                                                   data_list, table_header))
        invalid_option = "test_col"
        self.assertEqual("", self.host_info.get_option_value("IP", [], table_header))
        self.assertEqual("", self.host_info.get_option_value(invalid_option, data_list, table_header))
        self.assertTrue(f"Please check whether the {invalid_option} is in the table" in self.host_info.error_msg)

    def test_process_check_host(self):
        self.host_info.process_check_host()

    def test_is_valid_ip(self):
        # ipv4
        self.assertTrue(self.host_info.is_valid_ip("10.10.10.10"))
        # ipv6
        self.assertTrue(self.host_info.is_valid_ip("fec0:0090:1c02::3101"))
        # invalid
        self.assertFalse(self.host_info.is_valid_ip("10.x.10.10"))
        self.assertFalse(self.host_info.is_valid_ip(""))

    def test_check_ip(self):
        # ipv4
        self.host_info.ip = "10.10.10.10"
        self.host_info.check_ip()
        self.assertTrue(self.host_info.error_msg == [])
        # ipv6
        self.host_info.ip = "fec0:0090:1c02::3101"
        self.host_info.check_ip()
        self.assertTrue(self.host_info.error_msg == [])
        # invalid
        self.host_info.ip = ".10.10"
        self.host_info.check_ip()
        self.assertTrue("Invalid IP: .10.10" in self.host_info.error_msg)
        # empty
        self.host_info.ip = ""
        self.host_info.check_ip()
        self.assertTrue("Please fill in IP" in self.host_info.error_msg)

    def test_check_ansible_ssh_user(self):
        # valid
        self.host_info.ansible_ssh_user = "root"
        self.host_info.check_ansible_ssh_user()
        self.assertTrue(self.host_info.error_msg == [])

        # invalid
        self.host_info.ansible_ssh_user = "user"
        self.host_info.check_ansible_ssh_user()
        self.assertTrue("Invalid ansible_ssh_user: user" in self.host_info.error_msg)

        # empty
        self.host_info.ansible_ssh_user = ""
        self.host_info.check_ansible_ssh_user()
        self.assertTrue("Please fill in ansible_ssh_user" in self.host_info.error_msg)


class TestHccnHostInfo(unittest.TestCase):
    def setUp(self):
        table_header = ["Title【勿动】", "IP", "ansible_ssh_user", "ansible_ssh_pass", "ansible_ssh_port", "set_hostname",
                        "npu_num", "davinci", "NPU0", "NPU1", "NPU2","NPU3","NPU4",	"NPU5", "NPU6", "NPU7", "NPU8",
                        "NPU9", "NPU10", "NPU11", "NPU12", "NPU13", "NPU14", "NPU15", "NPU0", "NPU1", "NPU2", "NPU3",
                        "NPU4", "NPU5", "NPU6", "NPU7", "NPU8", "NPU9", "NPU10", "NPU11", "NPU12", "NPU13", "NPU14",
                        "NPU15", "gateways", "netmask", "bitmap", "dscp_tc", "common_network", "roce_port", "SCALE",
                        "RUNNER_IP", "WEIGHTS_PATH"]
        self.hccn_host_info = HccnHostInfo([], [], table_header)

    def test_process_check_host(self):
        self.hccn_host_info.process_check_host()

    def test_process_check_var(self):
        self.hccn_host_info.process_check_var()

    def test_get_host_info(self):
        self.hccn_host_info.ip = "2.2.2.2"
        self.hccn_host_info.ansible_ssh_user = "root"
        self.hccn_host_info.ansible_ssh_pass = ""
        self.hccn_host_info.device_ip = []
        self.hccn_host_info.detect_ip = []
        line = f'{self.hccn_host_info.ip} ansible_ssh_user="{self.hccn_host_info.ansible_ssh_user}" deviceip= detectip='
        self.assertEqual(line, self.hccn_host_info.get_host_info())

    def test_get_var_info(self):
        self.hccn_host_info.gateways = ""
        self.hccn_host_info.netmask = ""
        self.hccn_host_info.roce_port = "4791"
        self.hccn_host_info.bitmap = ""
        self.hccn_host_info.dscp_tc = ""
        self.hccn_host_info.common_network = ""
        line = ('gateways=""\n'
                'netmask=""\n'
                'roce_port="4791"\n'
                'bitmap=""\n'
                'dscp_tc=""\n'
                'common_network=""\n')
        self.assertEqual(line, self.hccn_host_info.get_var_info())

    def test_is_ipv6(self):
        self.assertTrue(self.hccn_host_info._is_ipv6("fec0:0090:1c02::3101"))
        self.assertFalse(self.hccn_host_info._is_ipv6("10.10.10.10"))

    def test__check_npu_ip(self):
        self.hccn_host_info.is_ipv6 = False
        empty_list = ["", "", "", "", "", "", "", "", "", "", "", "", "", "", "", ""]
        self.hccn_host_info._check_npu_ip(empty_list)
        self.hccn_host_info._check_npu_ip(["1.1.1.1", "1.1.1.2", "1.1.1.3", "1.1.1.4", "1.1.1.5", "1.1.1.6", "1.1.1.777",
                                           "fec0:0090:1c02::3101", "", "", "", "", "", "", "", ""])
        err_msg = [f"Please confirm the number of npu ip is correct.The number of npu ip should be 8 or 16.The npu ip is"
                   f" {empty_list}",
                   "The npu ip fec0:0090:1c02::3101 is different from deviceip[0]'s Internet Protocol version, deviceip,"
                   " detectip, gateway and netmask all should be set to IPv4 or IPv6.",
                   "Invalid npu IP address: 1.1.1.777"]
        for msg in err_msg:
            self.assertTrue(msg in self.hccn_host_info.error_msg)

    def test_check_device_ip(self):
        self.hccn_host_info.check_device_ip()

    def test_check_detect_ip(self):
        self.hccn_host_info.check_detect_ip()

    def test_check_bitmap(self):
        self.hccn_host_info.bitmap = ""
        self.hccn_host_info._check_bitmap()
        self.assertTrue(self.hccn_host_info.error_msg == [])

        self.hccn_host_info.bitmap = "0,0,0,0,1,0,0,0"
        self.hccn_host_info._check_bitmap()
        self.assertTrue(self.hccn_host_info.error_msg == [])

        self.hccn_host_info.bitmap = "0,2,1,0"
        self.hccn_host_info._check_bitmap()
        self.assertTrue("Invalid bitmap: {}.The bitmap you configured consist of 0 and 1, combined by ',', and the "
                        "length should be 8".format(self.hccn_host_info.bitmap) in self.hccn_host_info.error_msg)

    def test_check_dscp_tc(self):
        self.hccn_host_info.dscp_tc = ""
        self.hccn_host_info._check_dscp_tc()
        self.assertTrue(self.hccn_host_info.error_msg == [])

        self.hccn_host_info.dscp_tc = "33:2,"
        self.hccn_host_info._check_dscp_tc()
        self.assertTrue(self.hccn_host_info.error_msg == [])

        self.hccn_host_info.dscp_tc = "73:2,"
        self.hccn_host_info._check_dscp_tc()
        self.assertTrue(f"Invalid dscp_tc: {self.hccn_host_info.dscp_tc}" in self.hccn_host_info.error_msg)

    def test_check_common_network(self):
        self.hccn_host_info.common_network = ""
        self.hccn_host_info._check_common_network()
        self.assertTrue(self.hccn_host_info.error_msg == [])

        self.hccn_host_info.common_network = "0.0.0.0/0"
        self.hccn_host_info._check_common_network()
        self.assertTrue(self.hccn_host_info.error_msg == [])

        self.hccn_host_info.common_network = "1.0.0.0/0"
        self.hccn_host_info._check_common_network()
        self.assertTrue(f"Invalid common_network: {self.hccn_host_info.common_network}" in self.hccn_host_info.error_msg)

    def test_check_netmask(self):
        self.hccn_host_info.netmask = ""
        self.hccn_host_info._check_netmask()
        self.assertTrue("Please fill in netmask" in self.hccn_host_info.error_msg)

        self.hccn_host_info.netmask = "255.255.255.0"
        self.hccn_host_info.is_ipv6 = True
        self.hccn_host_info._check_netmask()
        self.assertTrue(f"The netmask {self.hccn_host_info.netmask} is different from deviceip[0]'s Internet Protocol version, "
                        f"deviceip, detectip, gateway and netmask all should be set to IPv4 or IPv6."
                        in self.hccn_host_info.error_msg)

        self.hccn_host_info.netmask = "112"
        self.hccn_host_info.is_ipv6 = False
        self.hccn_host_info._check_netmask()
        self.assertTrue(f"The netmask {self.hccn_host_info.netmask} is different from deviceip[0]'s Internet Protocol version, "
                        f"deviceip, detectip, gateway and netmask all should be set to IPv4 or IPv6."
                        in self.hccn_host_info.error_msg)

    def test_check_gateways(self):
        self.hccn_host_info.gateways = ""
        self.hccn_host_info._check_gateways()
        self.assertTrue("Please fill in gateways" in self.hccn_host_info.error_msg)

        self.hccn_host_info.gateways = "10.20.0.1"
        self.hccn_host_info.is_ipv6 = True
        self.hccn_host_info._check_gateways()
        self.assertTrue(f"The gateway {self.hccn_host_info.gateways} is different from deviceip[0]'s Internet Protocol "
                        f"version, deviceip, detectip, gateway and netmask all should be set to IPv4 or IPv6."
                        in self.hccn_host_info.error_msg)

        self.hccn_host_info.gateways = "1000.20.0.1"
        self.hccn_host_info._check_gateways()
        self.assertTrue(f"Invalid gateways: {self.hccn_host_info.gateways}" in self.hccn_host_info.error_msg)

        self.hccn_host_info.gateways = "10.20.0.1"
        self.hccn_host_info.device_ip = ["100.20.0.2"]
        self.hccn_host_info.netmask = "255.255.255.0"
        self.hccn_host_info._check_gateways()
        self.assertTrue("Please check the network configuration. The IP：100.20.0.2 are not in the same subnet."
                        in self.hccn_host_info.error_msg)

    def test_ip_to_binary(self):
        self.assertEqual("00000001000000010000001000000000", self.hccn_host_info._ip_to_binary("1.1.2.0"))
        self.assertEqual("11111110110000000000000010010000000111000000001000000000000000000000000000000000000000000"
                         "000000000000000000000000011000100000001",
                         self.hccn_host_info._ip_to_binary("fec0:0090:1c02::3101"))
        self.assertEqual(None, self.hccn_host_info._ip_to_binary("test"))

    def test_in_same_subnet(self):
        self.hccn_host_info.netmask = "255.255.255.0"
        self.assertTrue(self.hccn_host_info._in_same_subnet("10.1.1.2", "10.1.1.3"))
        self.assertFalse(self.hccn_host_info._in_same_subnet("10.1.20.2", "10.1.1.3"))
        self.assertFalse(self.hccn_host_info._in_same_subnet("10.1.20.2", "fec0:0090:1c02::3101"))
        self.assertTrue("Please fill in both the IP:10.1.20.2 and gateway:fec0:0090:1c02::3101 in either IPv4 or IPv6."
                         in self.hccn_host_info.error_msg)
        self.hccn_host_info.netmask = ""
        self.assertTrue(self.hccn_host_info._in_same_subnet("101.1.20.2", "10.1.1.3"))


class TestCommonHostInfo(unittest.TestCase):
    def setUp(self):
        table_header = ["Title【勿动】", "IP", "ansible_ssh_user", "ansible_ssh_pass", "ansible_ssh_port", "set_hostname",
                        "npu_num", "davinci", "NPU0", "NPU1", "NPU2","NPU3","NPU4",	"NPU5", "NPU6", "NPU7", "NPU8",
                        "NPU9", "NPU10", "NPU11", "NPU12", "NPU13", "NPU14", "NPU15", "NPU0", "NPU1", "NPU2", "NPU3",
                        "NPU4", "NPU5", "NPU6", "NPU7", "NPU8", "NPU9", "NPU10", "NPU11", "NPU12", "NPU13", "NPU14",
                        "NPU15", "gateways", "netmask", "bitmap", "dscp_tc", "common_network", "roce_port", "SCALE",
                        "RUNNER_IP", "WEIGHTS_PATH"]
        self.comm_host_info = CommonHostInfo([], [], table_header)

    def test_process_check_host(self):
        self.comm_host_info.process_check_host()

    def test_process_check_var(self):
        self.comm_host_info.process_check_var()

    def test_get_host_info(self):
        self.comm_host_info.ip = "2.2.2.2"
        self.comm_host_info.ansible_ssh_user = "root"
        self.comm_host_info.ansible_ssh_pass = "test"
        self.comm_host_info.hostname = "worker-1"
        self.comm_host_info.npu_num = "8"
        self.comm_host_info.davinci = "0,1,2"
        line = '2.2.2.2 ansible_ssh_user="root" ansible_ssh_pass="test" set_hostname="worker-1" npu_num=8 davinci=0,1,2 '
        self.assertEqual(line, self.comm_host_info.get_host_info())

    def test_get_var_info(self):
        self.comm_host_info.scale = "true"
        self.comm_host_info.runner_ip = ""
        self.comm_host_info.weights_path = "/test"
        line = ('SCALE="true"\n'
                'RUNNER_IP=""\n'
                'WEIGHTS_PATH="/test"\n')
        self.assertEqual(line, self.comm_host_info.get_var_info())

    def test_check_ansible_ssh_port(self):
        self.comm_host_info.ansible_ssh_port = "22"
        self.comm_host_info._check_ansible_ssh_port()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.ansible_ssh_port = ""
        self.comm_host_info._check_ansible_ssh_port()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.ansible_ssh_port = "aa"
        self.comm_host_info._check_ansible_ssh_port()
        self.assertTrue("Invalid ansible_ssh_port: aa" in self.comm_host_info.error_msg)

    def test_check_hostname(self):
        self.comm_host_info.hostname = "worker-1"
        self.comm_host_info._check_hostname()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.hostname = ""
        self.comm_host_info._check_hostname()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.hostname = "Test"
        self.comm_host_info._check_hostname()
        self.assertTrue("Invalid set_hostname: Test" in self.comm_host_info.error_msg)

    def test_check_npu_num(self):
        self.comm_host_info.npu_num = "8"
        self.comm_host_info._check_npu_num()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.npu_num = ""
        self.comm_host_info._check_npu_num()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.npu_num = "test"
        self.comm_host_info._check_npu_num()
        self.assertTrue("Invalid npu_num: test" in self.comm_host_info.error_msg)

    def test_check_davinci(self):
        self.comm_host_info.davinci = "0,1,2"
        self.comm_host_info._check_davinci()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.davinci = "0,1,1"
        self.comm_host_info._check_davinci()
        self.assertTrue("Invalid davinci: 0,1,1" in self.comm_host_info.error_msg)

        self.comm_host_info.npu_num = "4"
        self.comm_host_info.davinci = "0,1,2,3,4"
        self.comm_host_info._check_davinci()
        self.assertTrue(f"The npu_num must be greater than or equal to the number of davinci.The number of davinci is "
                        f"{len(self.comm_host_info.davinci.split(','))}. The npu_num is "
                        f"{int(self.comm_host_info.npu_num)}" in self.comm_host_info.error_msg)

    def test_check_scale(self):
        self.comm_host_info.scale = "True"
        self.comm_host_info._check_scale()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.scale = ""
        self.comm_host_info._check_scale()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.scale = "test"
        self.comm_host_info._check_scale()
        self.assertTrue("Invalid SCALE: test" in self.comm_host_info.error_msg)

    def test_check_runner_ip(self):
        # ipv4
        self.comm_host_info.runner_ip = "10.10.10.10"
        self.comm_host_info._check_runner_ip()
        self.assertTrue(self.comm_host_info.error_msg == [])
        # ipv6
        self.comm_host_info.runner_ip = "fec0:0090:1c02::3101"
        self.comm_host_info._check_runner_ip()
        self.assertTrue(self.comm_host_info.error_msg == [])
        # empty
        self.comm_host_info.runner_ip = ""
        self.comm_host_info._check_runner_ip()
        self.assertTrue(self.comm_host_info.error_msg == [])
        # invalid
        self.comm_host_info.runner_ip = "10.x.10.10"
        self.comm_host_info._check_runner_ip()
        self.assertTrue("Invalid RUNNER_IP: 10.x.10.10" in self.comm_host_info.error_msg)

    def test_check_weights_path(self):
        self.comm_host_info.weights_path = "/home"
        self.comm_host_info._check_weights_path()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.weights_path = "~/test"
        self.comm_host_info._check_weights_path()
        self.assertTrue(self.comm_host_info.error_msg == [])

        self.comm_host_info.weights_path = "test"
        self.comm_host_info._check_weights_path()
        self.assertTrue("Invalid WEIGHTS_PATH: test" in self.comm_host_info.error_msg)

        self.comm_host_info.davinci = "0,1"
        self.comm_host_info.weights_path = ""
        self.comm_host_info._check_weights_path()
        self.assertTrue("Please fill in WEIGHTS_PATH" in self.comm_host_info.error_msg)


class TestInventoryInfo(unittest.TestCase):
    def setUp(self):
        self.inventory_info = InventoryInfo()

    def test_trans_to_inventory_info(self):
        data_list = [["worker", "1.1.1.1", "root", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "",
                      "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "10.20.0.1,10.20.1.1", "112",
                     "0,0,0,0,0,0,1,0", "33:2,", "0.0.0.0/0", "", "true", "1.1.1.1", "/home/mindie"], []]
        self.inventory_info.table_header = ["Title【勿动】", "IP", "ansible_ssh_user", "ansible_ssh_pass", "ansible_ssh_port", "set_hostname",
                        "npu_num", "davinci", "NPU0", "NPU1", "NPU2","NPU3","NPU4",	"NPU5", "NPU6", "NPU7", "NPU8",
                        "NPU9", "NPU10", "NPU11", "NPU12", "NPU13", "NPU14", "NPU15", "NPU0", "NPU1", "NPU2", "NPU3",
                        "NPU4", "NPU5", "NPU6", "NPU7", "NPU8", "NPU9", "NPU10", "NPU11", "NPU12", "NPU13", "NPU14",
                        "NPU15", "gateways", "netmask", "bitmap", "dscp_tc", "common_network", "roce_port", "SCALE",
                        "RUNNER_IP", "WEIGHTS_PATH"]
        self.inventory_info.trans_to_inventory_info(data_list, data_list[0])
        self.assertTrue(self.inventory_info.error_msg == [])
        self.assertTrue(self.inventory_info.hosts_info[self.inventory_info.MASTER_SEC] == [])
        self.assertTrue(self.inventory_info.hosts_info[self.inventory_info.HCCN_SEC] == [])
        self.assertTrue(self.inventory_info.hosts_info[self.inventory_info.WORKER_SEC] != [])

        data_list = [["worker", "1.1.1.1.1", "root", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "",
                      "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "10.20.0.1,10.20.1.1", "112",
                     "0,0,0,0,0,0,1,0", "33:2,", "0.0.0.0/0", "", "true", "1.1.1.1", "/home/mindie"], []]
        with self.assertRaises(Exception):
            self.inventory_info.trans_to_inventory_info(data_list, data_list[0])
        self.assertTrue("Invalid IP: 1.1.1.1.1" in self.inventory_info.error_msg)

